package com.chenjt.five_embedding;

import com.chenjt.three_tools.ModelUtil;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

/**
 * 注释：
 *
 * @author chenjt
 * @date 2025/3/10 11:26
 */
public class VectorDemo {

    public static void main(String[] args) {

        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey(ModelUtil.API_KEY)
                .baseUrl(ModelUtil.BASE_URL)
                .build();

        // 查看embedding结果
//        getEmbeddindResult(embeddingModel);


        // 将embedding结果存储到redis中
//        embeddingResultToRedis(embeddingModel);

        // 比较embedding向量结果的相似度
        compareEmbeddingResult(embeddingModel);


    }

    /**
     * 比较embedding向量结果的相似度，相似度是从0-1的数字，越接近1，越相似
     * @param embeddingModel
     */
    private static void compareEmbeddingResult(EmbeddingModel embeddingModel) {
        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536)
                .build();

        /**
         * 去redis中去找和"我的名字叫陈橙晨"这句话相关的内容
         * 第2个参数：最多查3条
         * 第3个参数：相似度下限-1
         */
        List<EmbeddingMatch<TextSegment>> matches
                = embeddingStore.findRelevant(embeddingModel.embed("我的名字叫陈橙晨").content(), 3, -1);
        // redis中存了“你好，我是陈橙晨”这句话的向量结果数据，和这句话的相似度
        for (EmbeddingMatch<TextSegment> match: matches){
            // 0.97277015447615，相似度是从0-1的数字，越接近1，越相似
            System.out.println(match.score());
        }
    }

    /**
     * 将embedding结果存入redis中
     * （
     *  redis需要装stack，这里没有装，直接用docker去跑
     *  docker run -d -p 6379:6379 redis/redis-stack-server:latest
     * ）
     * @param embeddingModel
     */
    private static void embeddingResultToRedis(EmbeddingModel embeddingModel) {
        Response<Embedding> embed = getEmbeddindResult(embeddingModel);

        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536)
                .build();
        embeddingStore.add(embed.content());
    }

    /**
     * 查看embedding的结果是怎么样的
     * @param embeddingModel
     */
    private static Response<Embedding> getEmbeddindResult(EmbeddingModel embeddingModel) {
        // 类似于一个二位的维度，在平面坐标轴上表示（x,y），只是这里向量出来的纬度有1000多个维度
        Response<Embedding> embed = embeddingModel.embed("你好，我是陈橙晨");
        // 查看向量的结果
        System.out.println(embed.content().toString());
        // 查看向量结果的维度(OpenAiEmbeddingModel的维度是1536)
        System.out.println(embed.content().vector().length);
        return embed;
    }






}
